package com.yhq.thrift.client;

import lombok.extern.slf4j.Slf4j;
import org.apache.thrift.protocol.TBinaryProtocol;
import org.apache.thrift.protocol.TMultiplexedProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TTransport;

import java.text.MessageFormat;
import java.util.List;
import java.util.Map;
import java.util.Random;

/**
 * Created by yanghuqiang(yhqnh@aliyun.com) on 2016/8/15.
 */
@Slf4j
public class ThriftConnectionFactory {

    private static ThriftGenericKeyedObjectPool<ServerKey, TTransport> pool;

    private ThriftPoolConfig poolConfig = new ThriftPoolConfig();

    private Map<String, ServerInfo> serverInfoMap;

    public ThriftConnectionFactory() throws ThriftClientException {

    }

    public ThriftConnectionFactory(ThriftPoolConfig config) throws ThriftClientException {
        this.poolConfig = poolConfig != null ? poolConfig : new ThriftPoolConfig();
    }

    public void createPool() {
        if (poolConfig != null) {
            ThriftPoolFactory factory = new ThriftPoolFactory();
            this.pool = new ThriftGenericKeyedObjectPool<ServerKey, TTransport>(factory, this.poolConfig);
        }
    }

    public ThriftConnection getConnection(String serviceName, String serviceClient) throws Exception {
        ServerKey serverKey = getLoadBanlance(serviceName);
        TTransport transport = pool.borrowObject(serverKey);
        Object client = getThriftClient(serviceName, serviceClient, transport);
        ThriftConnection connection = new ThriftConnection(client, serverKey, transport, pool);
        return connection;
    }

    public Object getThriftClient(String serviceName, String serviceClient, TTransport transport) throws Exception {
        TBinaryProtocol protocol = new TBinaryProtocol(transport);
        TMultiplexedProtocol mp = new TMultiplexedProtocol(protocol, serviceName);
        Object object = Class.forName(serviceClient).getConstructor(TProtocol.class).newInstance(mp);
        return object;
    }

    public ServerKey getLoadBanlance(String serviceName) {
        ServerInfo serverInfo = serverInfoMap.get(serviceName);
        List<ServerKey> serverKeyList = serverInfo.getServerKeyList();
        Random random = new Random();
        int index = random.nextInt(serverKeyList.size());
        ServerKey serverKey = serverKeyList.get(index);
        if (serverKey == null) {
            log.info("没有可用的服务{}", serviceName);
            throw new ThriftClientException(MessageFormat.format("没有可用的服务{0}", serviceName));
        }
        return serverKey;
    }

    public void destroy() {
        if (pool != null) {
            pool.close();
        }
    }

    public Map<String, ServerInfo> getServerInfoMap() {
        return serverInfoMap;
    }

    public void setServerInfoMap(Map<String, ServerInfo> serverInfoMap) {
        this.serverInfoMap = serverInfoMap;
    }

    public ThriftPoolConfig getPoolConfig() {
        return poolConfig;
    }

    public void setPoolConfig(ThriftPoolConfig poolConfig) {
        this.poolConfig = poolConfig;
    }
}
